import pymongo
from config import br_mongo_conf
from base.tools import wrapper
from contextlib import contextmanager
from base.logcontrol import logger
import sys


class MongoBase(object):

    def __init__(self, database):
        self.database = database

    def __repr__(self):
        return "MongoBase对象"

    @contextmanager
    def connector(self):
        """
        设置查询的数据库
        :return:
        """
        conn = pymongo.MongoClient("mongodb://{user}:{password}@{host}/?authSource=admin".format(**br_mongo_conf))
        try:
            yield conn.get_database(self.database)
        except Exception as e:
            logger.error(f"\n\t链接数据出错\n"
                         f"\t异常信息：{e}\n")
            sys.exit()
        finally:
            conn.close()

    @wrapper
    def get_data(self, collection: str, sql: dict, filters: dict = None, lt: int = 100):
        """
        执行查询
        :param collection: 查询的集合
        :param sql: 查询语句
        :param filters: 过滤条件
        :param lt: 限制条数
        :return:
        """
        _data = []
        with self.connector() as db:
            col = db[collection]
            obj = col.find(sql, filters, limit=lt)
            for elm in obj:
                _data.append(elm) if elm else None
        return _data
